// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/table_builder.h>

#include <memory>
#include <utility>

#include <nebula/array/array_base.h>
#include <nebula/array/builder_base.h>
#include <nebula/core/record_batch.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/log/logging.h>

namespace nebula {

    // ----------------------------------------------------------------------
    // RecordBatchBuilder

    RecordBatchBuilder::RecordBatchBuilder(const std::shared_ptr<Schema> &schema,
                                           MemoryPool *pool, int64_t initial_capacity)
            : schema_(schema), initial_capacity_(initial_capacity), pool_(pool) {}

    turbo::Result<std::unique_ptr<RecordBatchBuilder>> RecordBatchBuilder::create(
            const std::shared_ptr<Schema> &schema, MemoryPool *pool) {
        return create(schema, pool, kMinBuilderCapacity);
    }

    turbo::Result<std::unique_ptr<RecordBatchBuilder>> RecordBatchBuilder::create(
            const std::shared_ptr<Schema> &schema, MemoryPool *pool, int64_t initial_capacity) {
        auto builder = std::unique_ptr<RecordBatchBuilder>(
                new RecordBatchBuilder(schema, pool, initial_capacity));
        TURBO_RETURN_NOT_OK(builder->CreateBuilders());
        TURBO_RETURN_NOT_OK(builder->InitBuilders());
        return builder;
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatchBuilder::flush(bool reset_builders) {
        std::vector<std::shared_ptr<Array>> fields;
        fields.resize(this->num_fields());

        int64_t length = 0;
        for (int i = 0; i < this->num_fields(); ++i) {
            TURBO_RETURN_NOT_OK(raw_field_builders_[i]->finish(&fields[i]));
            if (i > 0 && fields[i]->length() != length) {
                return turbo::invalid_argument_error("All fields must be same length when calling flush");
            }
            length = fields[i]->length();
        }

        // For certain types like dictionaries, types may not be fully
        // determined before we have flushed. Make sure that the RecordBatch
        // gets the correct types in schema.
        // See: #ARROW-9969
        std::vector<std::shared_ptr<Field>> schema_fields(schema_->fields());
        for (int i = 0; i < this->num_fields(); ++i) {
            if (!schema_fields[i]->type()->equals(fields[i]->type())) {
                schema_fields[i] = schema_fields[i]->with_type(fields[i]->type());
            }
        }
        std::shared_ptr<Schema> schema =
                std::make_shared<Schema>(std::move(schema_fields), schema_->metadata());

        std::shared_ptr<RecordBatch> batch =
                RecordBatch::create(std::move(schema), length, std::move(fields));

        if (reset_builders) {
            TURBO_RETURN_NOT_OK(InitBuilders());
        }

        return batch;
    }

    turbo::Result<std::shared_ptr<RecordBatch>> RecordBatchBuilder::flush() { return flush(true); }

    void RecordBatchBuilder::set_initial_capacity(int64_t capacity) {
        KCHECK_GT(capacity, 0) << "Initial capacity must be positive";
        initial_capacity_ = capacity;
    }

    turbo::Status RecordBatchBuilder::CreateBuilders() {
        field_builders_.resize(this->num_fields());
        raw_field_builders_.resize(this->num_fields());
        for (int i = 0; i < this->num_fields(); ++i) {
            TURBO_RETURN_NOT_OK(MakeBuilder(pool_, schema_->field(i)->type(), &field_builders_[i]));
            raw_field_builders_[i] = field_builders_[i].get();
        }
        return turbo::OkStatus();
    }

    turbo::Status RecordBatchBuilder::InitBuilders() {
        for (int i = 0; i < this->num_fields(); ++i) {
            TURBO_RETURN_NOT_OK(raw_field_builders_[i]->Reserve(initial_capacity_));
        }
        return turbo::OkStatus();
    }

}  // namespace nebula
